#include "../../../../include/util.h"
#include "sync-func.h"

#include "llvm/DerivedTypes.h"
#include "llvm/Instructions.h"
#include "llvm/Constants.h"
#include "llvm/BasicBlock.h"
#include "llvm/Pass.h"
#include "llvm/Support/CFG.h"
#include "llvm/Transforms/Utils/BasicInliner.h"
#include "llvm/Target/TargetData.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
#include "llvm/ADT/DenseSet.h"

#include <set>

#include "../../include/tern/TernDefs.h"
//#include "../../include/tern/TernLandmark.h"

extern const char* landmarks[TERN_LANDMARKS_NUM];
extern bool isIgnoredLandmarkInReplay(unsigned index);

using namespace llvm;
using namespace std;

namespace defens {

	bool SyncFunc::is_exit_block(BasicBlock *bb) {
		if (succ_begin(bb) != succ_end(bb))
			return false;
		for (BasicBlock::iterator ii = bb->begin(); ii != bb->end(); ii++) {
			if (dyn_cast<ReturnInst>(ii))
				return true;
		}
		return false;
	}

	void SyncFunc::construct_call_graph(Module &M) {
		/* Which functions are used as function pointers? */
		set<Function *> fps;
		forallinst(M, ii) {
			if (is_intrinsic_call(ii))
				continue;
			unsigned start_idx = 0;
			if (is_call(ii))
				start_idx = 1;
			unsigned n = ii->getNumOperands();
			for (unsigned i = start_idx; i < n; ++i) {
				Value *v = ii->getOperand(i);
				if (isa<Function>(v)) {
					Function *f = dyn_cast<Function>(v);
					fps.insert(f);
				}
			}
		}
#if 1
		cerr << "# of functions used as function pointers = " << fps.size() << endl;
		for (set<Function *>::iterator it = fps.begin(); it != fps.end(); ++it)
			cerr << (*it)->getNameStr() << endl;
#endif
		/* Every BB's called functions */
		called_func.clear();
		call_sites.clear();
		vector<int> dist_n_cands;
		forallinst(M, ii) {
			if (is_call(ii) && !is_intrinsic_call(ii)) {
				CallSite cs(ii);
				Function *callee = cs.getCalledFunction();
				if (!callee) {
					Value *v = cs.getCalledValue();
					assert(v);
					assert(isa<PointerType>(v->getType()));
					const Type *t = dyn_cast<PointerType>(v->getType())->getElementType();
					assert(isa<FunctionType>(t));
#if 0
					cerr << "Function Pointer detected\n";
					cerr << "Candidates:\n";
#endif
					int n_cands = 0;
					for (set<Function *>::iterator it = fps.begin(); it != fps.end(); ++it) {
						Function *fi = *it;
						if (fi->getFunctionType() == t) {
							// cerr << fi->getNameStr() << ' ';
							called_func[ii].push_back(fi);
							call_sites[fi].push_back(ii);
							n_cands++;
						}
					}
					dist_n_cands.push_back(n_cands);
					// cerr << endl;
				} else {
					called_func[ii].push_back(callee);
					call_sites[callee].push_back(ii);
				}
			}
		}
#if 0
		cerr << "Distribution of # of candidates:\n";
		sort(dist_n_cands.begin(), dist_n_cands.end());
		for (size_t i = 0; i < dist_n_cands.size(); ++i) {
			cerr << dist_n_cands[i] << ' ';
		}
		cerr << endl;
#endif
		for (DenseMap<Instruction *, vector<Function *> >::iterator it = called_func.begin();
				it != called_func.end(); ++it) {
			sort(it->second.begin(), it->second.end());
			it->second.resize(unique(it->second.begin(), it->second.end()) - it->second.begin());
		}
		for (DenseMap<Function *, vector<Instruction *> >::iterator it = call_sites.begin();
				it != call_sites.end(); ++it) {
			sort(it->second.begin(), it->second.end());
			it->second.resize(unique(it->second.begin(), it->second.end()) - it->second.begin());
		}
		exits.clear();
		forallbb(M, bi) {
			if (is_exit_block(bi))
				exits[bi->getParent()].push_back(&bi->back());
		}
		fprintf(stderr, "Finished generating the call graph.\n");
	}

	void SyncFunc::setup_sync_functions(Module &M) {
		Function *f;
		sync_funcs.clear();
		for (unsigned i = 0; i < TERN_LANDMARKS_NUM; i++) {
			if (!isIgnoredLandmarkInReplay(i)) {
				if ((f = M.getFunction(landmarks[i])) != NULL) {
					sync_funcs.push_back(f);
					fprintf(stderr, "Inserted landmark %s for condition slicing.\n", landmarks[i]);
				}
			}
		}
		/*f = M.getFunction("pthread_mutex_lock");
		if (f)
			sync_funcs.push_back(f);
		f = M.getFunction("pthread_mutex_unlock");
		if (f)
			sync_funcs.push_back(f);
		f = M.getFunction("pthread_create");
		if (f)
			sync_funcs.push_back(f);
		f = M.getFunction("pthread_join");
		if (f)
			sync_funcs.push_back(f);*/
	}

	// TODO: search in vector may be slow
	bool SyncFunc::is_sync_function(Function *f) {
		for (size_t i = 0; i < sync_funcs.size(); ++i) {
			if (sync_funcs[i] == f)
				return true;
		}
		return false;
	}

	vector<Function *> SyncFunc::get_called_functions(Instruction *ins) {
		DenseMap<Instruction *, vector<Function *> >::iterator it = called_func.find(ins);
		if (it == called_func.end()) {
			//fprintf(stderr, "[Warning] <ins> is not in the call graph.\n");
			return vector<Function *>();
		}
		return it->second;
	}

#if 0
	bool SyncFunc::no_lock(const Function *f) {
		string name = f->getNameStr();
		if (name == "apr_palloc") return true;
		if (name == "ap_log_error") return true;
		return false;
	}
#endif

	void SyncFunc::DFS(Function *f) {
#if 0
		if (no_lock(f))
			return;
#endif
		visited.insert(f);
		vector<Instruction *> &css = call_sites[f];
		for (size_t i = 0; i < css.size(); ++i) {
			Function *caller = css[i]->getParent()->getParent();
			if (visited.count(caller) == 0) {
				parent[caller] = f;
				DFS(caller);
			}
		}
	}

	void SyncFunc::output(const Module &M) const {
		vector<string> func_names;
		FILE *fout;

		func_names.clear();
		for (Module::const_iterator fi = M.begin(); fi != M.end(); ++fi)
			func_names.push_back(fi->getNameStr());
		sort(func_names.begin(), func_names.end());
		fout = fopen("/tmp/all-func.txt", "w");
		for (size_t i = 0; i < func_names.size(); ++i) {
			fprintf(fout, "%s\n", func_names[i].c_str());
		}

		func_names.clear();
		for (DenseSet<Function *>::const_iterator it = visited.begin();
				it != visited.end(); ++it) {
			func_names.push_back((*it)->getNameStr());
		}
		sort(func_names.begin(), func_names.end());
		fout = fopen("/tmp/sync-func.txt", "w");
		for (size_t i = 0; i < func_names.size(); ++i) {
			fprintf(fout, "%s\n", func_names[i].c_str());
		}
	}

	bool SyncFunc::contains_sync_operation(Function *f) {
		return visited.count(f);
	}

	void SyncFunc::traverse_call_graph(Module &M) {
		visited.clear();
		parent.clear();
		for (size_t i = 0; i < sync_funcs.size(); ++i) {
			DFS(sync_funcs[i]);
		}
	}

	void SyncFunc::print_call_chain(Function *f) {
		if (!parent.count(f)) {
			cerr << endl;
			return;
		}
		cerr << f->getNameStr() << " => ";
		print_call_chain(parent[f]);
	}

	void SyncFunc::print_calling_functions(Function *f) {
		if (!call_sites.count(f))
			return;
		cerr << f->getNameStr() << ": ";
		vector<Instruction *> &callers = call_sites[f];
		for (size_t i = 0; i < callers.size(); ++i) {
			cerr << callers[i]->getParent()->getParent()->getNameStr() << ' ';
		}
		cerr << endl;
	}

	void SyncFunc::stats(const Module &M) const {
	}

	bool SyncFunc::runOnModule(Module &M) {
		construct_call_graph(M);
		setup_sync_functions(M);
		traverse_call_graph(M);
		return false;
	}

	void SyncFunc::print(llvm::raw_ostream &O, const Module *M) const {
		output(*M);
		stats(*M);
	}

	char SyncFunc::ID = 0;
}

namespace {
    static RegisterPass<defens::SyncFunc> X(
			"sync-func",
			"Get functions with sync operations",
			false,
			true); // is analysis
}
